package cn.heycm.cachedrequest.filter.wrapper;


import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

public class CachedRequestWrapper extends HttpServletRequestWrapper {

    private final byte[] cachedBody;

    private Map<String, String> extHeaders;

    public CachedRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        ServletInputStream is = request.getInputStream();
        cachedBody = is.readAllBytes();
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        return new CachedServletInputStream(cachedBody);
    }

    @Override
    public BufferedReader getReader() throws IOException {
        ByteArrayInputStream is = new ByteArrayInputStream(cachedBody);
        return new BufferedReader(new InputStreamReader(is));
    }

    public byte[] getContentAsByteArray() {
        return cachedBody;
    }

    public String getContentAsString() {
        return new String(cachedBody, StandardCharsets.UTF_8);
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        if (null == value && extHeaders != null) {
            value = extHeaders.get(name);
        }
        return value;
    }

    public void addHeader(String name, String value) {
        if (extHeaders == null) {
            extHeaders = new HashMap<>();
        }
        extHeaders.put(name, value);
    }


    private static class CachedServletInputStream extends ServletInputStream {

        private final ByteArrayInputStream is;

        private CachedServletInputStream(byte[] body) {
            this.is = new ByteArrayInputStream(body);
        }

        @Override
        public boolean isFinished() {
            return is.available() == 0;
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setReadListener(ReadListener readListener) {
            throw new UnsupportedOperationException();
        }

        @Override
        public int read() throws IOException {
            return is.read();
        }
    }
}
